#!/usr/bin/env python

import sys
import traceback
from hmm import *
from optparse import OptionParser

class TestFailed: pass

def check(b):
    """Like assert, except throws an exception we can actually catch"""
    if not b:
        raise TestFailed

def test_learn_from_labeled_data():
    state_seq = [[0,1,1,0,1,0,1,1], [0,0,1,0]]
    obs_seq =   [[0,0,1,1,0,0,0,1], [0,1,0,0]]
    hmm = HMM(range(2), range(2))
    hmm.learn_from_labeled_data(state_seq, obs_seq)
#    print hmm

    eps = 0.00001
    check(max_delta(hmm.initial, [0.750000,0.250000]) < eps)
    check(max_delta(hmm.transition, 
                     [[0.285714, 0.714286],
                      [0.571429, 0.428571]]) < eps)
    check(max_delta(hmm.observation,
                     [[0.625000, 0.375000],
                      [0.625000, 0.375000]]) < eps)


def simple_weather_model():
    hmm = HMM(['s1','s2'], ['R','NR'])
    init = [0.7, 0.3]
    trans = [[0.8,0.2],
             [0.1,0.9]]
    observ = [[0.75,0.25], 
              [0.4,0.6]]
    hmm.set_hidden_model(init, trans, observ)
    return hmm

def toy_model():
    hmm = HMM(['s1','s2'], ['R','NR'])
    init = [0.5, 0.5]
    trans = [[0.2,0.8],
             [0.8,0.2]]
    observ = [[0.8,0.2], 
              [0.2,0.8]]
    hmm.set_hidden_model(init, trans, observ)
    return hmm
    
def simple_viterbi_test():
    hmm = simple_weather_model()
    seq = [1, 1, 0]  # NR, NR, R

    hidden_seq = hmm.most_likely_states(seq)
#    print "most likely states for [NR, NR, R] = %s" % hidden_seq
    check(hidden_seq == [1,1,1])

def long_viterbi_test():
    hmm = toy_model()
    N = 10
    seq = [1,0,1,0,1,0,1,1,0] * 400
    hidden_seq = hmm.most_likely_states(seq, False)
    # Got right answer from the version with logs.
    check(hidden_seq[2000:2010] == [1, 0, 1, 0, 1, 1, 0, 1, 0, 1])



def get_model(hmm):
    return (hmm.initial, hmm.transition, hmm.observation)

def simple_weather_model_tests():
    """Some tests for the EM internal functions"""
    hmm = simple_weather_model()
#    print "HMM is:"
#    print hmm
    
    seq = [0,0]  # R, R
    alpha, logp = get_alpha(seq, get_model(hmm))
    beta = get_beta(seq, get_model(hmm))
    gamma = get_gamma(alpha, beta)
    xi = get_xi(seq, alpha, beta, get_model(hmm))

    eps = 0.0001
    check(max_delta(alpha, [[0.813953, 0.186047],
                            [0.791789,  0.208211]]) < eps)
    check(max_delta(beta, [[0.609865, 0.390135],
                           [0.5, 0.5]]) < eps)
    check(max_delta(gamma, [[  0.872434,  0.127566  ],
                            [  0.791789,  0.208211  ]]) < eps)
    check(max_delta(xi, [[[ 0.76979472,  0.1026393 ],
                          [ 0.02199413,  0.10557185]],
                         [[ 0., 0.],
                          [ 0., 0.]]]) < eps)

    p = False
    if p:
        print "alpha: "
        format_array_print (alpha)
        print "beta: "
        format_array_print (beta)
        print "gamma: "
        format_array_print (gamma)
        print "xi: "
        print xi

        p = exp(logp)
        print "prob ([R, R]): logp= %f  p= %f" % (logp, p)

def simple_weather_model_em_test():
    """ Not actually a test--just runs the code and prints things out.
    Should match the class notes.
    """
    # Test B-W
    import hmm
    hmm.PRODUCTION = False  # Don't want smoothing or normalization..

    m = simple_weather_model()
    seqs = [[0, 0], [1, 1, 0]]
    init_model = get_model(m)
    baumwelch(seqs, 2, 2, True, init_model)


def test_init_model():
    # Initialize things to specific values, for testing
    import hmm
    hmm.PRODUCTION = False
    N = 3
    # Normalized below
    transition = array([[1.0,1.0,1.0], 
                        [1.0,1.0,1.0], 
                        [1.0,1.0,1.0]])
    observation  = array([[1.0,1.0], [3.0,1.0], [1.0,3.0]])
    initial    = ones([N])

    # Normalize
    initial    = initial/sum(initial)
    for i in range(N): 
        transition[i,:] = transition[i,:]/sum(transition[i,:])
        observation[i,:]  = observation[i,:]/sum(observation[i,:])

    seq = [[0, 0, 0, 0, 1, 0, 1, 1, 1, 1]]

    # Run EM on this one sequence, with the initial model above.

    model = (initial, transition, observation)
    baumwelch(seq, 3, 2, True, model)


# Normal usage: without options, will run all these tests
tests = {'learn_from_labeled_data' : test_learn_from_labeled_data, 
         'simple_weather_model' : simple_weather_model_tests,
         'simple_viterbi' : simple_viterbi_test,
         'long_viterbi' : long_viterbi_test}
         


def main(argv):
    usage = "usage: %prog [options] (pass -h for more info)"
    parser = OptionParser(usage)

    parser.add_option("--basic-em",
                      action="store_true", dest="basic_em", default=False,
                      help="Run a simple EM test")

    (options, args) = parser.parse_args(argv[1:])

    if options.basic_em:
        test_init_model()
    else:
        for (k,f) in tests.items():
            try:
                print " =========== Running test: %s =========" % k
                f()
                print " >>>>>>  passed  >>>>"
            except:
                print "TEST FAILED: ", traceback.format_exc()
        print "========== DONE ============="

if __name__ == "__main__":
    main(sys.argv)
